import sqlparse
from sqlparse.sql import IdentifierList, Identifier, Function, Where
from sqlparse.tokens import Keyword, DML, Comparison, String, Number


class Parser:
    def query(self, sql: str):
        self.parsed = sqlparse.parse(sql)[0]

    def get_aggregate(self):
        for item in self.parsed:
            if isinstance(item, Function):
                yield item.get_name()
            if isinstance(item, IdentifierList):
                for il in item.tokens:
                    if isinstance(il, Function):
                        yield il.get_name()

    def get_predicates(self):
        for item in self.parsed:
            if isinstance(item, Where):
                for token in item.tokens:
                    if (
                        hasattr(token, "tokens")
                        and len(token.tokens) == 3
                        and token.tokens[1].ttype is Comparison
                    ):
                        if token.right.ttype is String.Single:
                            text = token.right.value
                            text = text.replace("'", "")
                            yield text
                        if token.right.ttype is Number.Integer:
                            yield float(token.right.value)

    def get_table_identifiers(self):
        for item in self.parsed:
            if isinstance(item, IdentifierList):
                for identifier in item.get_identifiers():
                    yield identifier.get_name()
            elif isinstance(item, Identifier):
                yield item.get_name()
            # It's a bug to check for Keyword here, but in the example
            # above some tables names are identified as keywords...
            elif item.ttype is Keyword:
                yield item.value

    def get_embedding_predicate(self):
        for item in self.parsed:
            if isinstance(item, Where):
                for token in item.tokens:
                    if (
                        hasattr(token, "tokens")
                        and len(token.tokens) == 3
                        and token.tokens[1].ttype is Comparison
                    ):
                        if token.right.ttype is String.Single:
                            text = token.right.value
                            text = text.replace("'", "")
                            yield text
                        if token.right.ttype is Number.Integer:
                            yield float(token.right.value)
            if isinstance(item, IdentifierList):
                for identifier in item.get_identifiers():
                    yield identifier.get_name()
            elif isinstance(item, Identifier):
                yield item.get_name()
            # It's a bug to check for Keyword here, but in the example
            # above some tables names are identified as keywords...
            elif item.ttype is Keyword:
                yield item.value


if __name__ == "__main__":
    parser = Parser()
    parser.query(
        "select count(*), avg(*) from foo where id=10 and age>20 and name='tom' "
    )
    # print(parser.parsed.tokens)
    print(list(parser.get_aggregate()))
    print(list(parser.get_predicates()))
    print(list(parser.get_embedding_predicate()))

    parser = Parser()
    parser.query("select count(*) from foo  ")
    # print(parser.parsed.tokens)
    print(list(parser.get_aggregate()))
    print(list(parser.get_predicates()))
    print(list(parser.get_embedding_predicate()))
